//
// Created by yunlang on 2021/12/1.
//

#include "HuffmanCompression.h"
#include<vector>
#include <unordered_map>
#include "hfmCodeBitSet.h"
#include<cstring>
using namespace std;

void HuffmanCompression::calcWeight(const unsigned char *dataPtr, unsigned int dataSize,
                                    unordered_map<unsigned char, unsigned int >& weightMap)
{
    for( unsigned int i = 0; i < dataSize; i++ )
    {
        if( weightMap.find(dataPtr[i]) == weightMap.end() )
            weightMap[dataPtr[i]] = 1;
        else
            weightMap[dataPtr[i]]++;
    }
}
void HuffmanCompression::generateEncodedNodeQueue(const unordered_map<unsigned char, unsigned int> &weightMap)
{
    for(const auto &ele : weightMap)
        nodeQueue.push( new TreeNode(ele.first,ele.second) );
}

void HuffmanCompression::generateTree()
{
    while( nodeQueue.getSize() > 1)
    {
        TreeNode * leftNode = nodeQueue.top();
        nodeQueue.pop();
        TreeNode * rightNode = nodeQueue.top();
        nodeQueue.pop();
        auto * parentNode = new TreeNode( 0,leftNode->weight + rightNode->weight);
        parentNode->left = leftNode;
        parentNode->right = rightNode;
        nodeQueue.push(parentNode);
    }
    treeRoot = nodeQueue.top();
}

void HuffmanCompression::dfs(TreeNode* node, hfmCodeBitSet& path,
                             unordered_map<unsigned char, hfmCodeBitSet> &resCodeMap)
{
    if(node->left == nullptr && node->right == nullptr)
    {
        resCodeMap[node->val] = path;
        return;
    }
    path.append(0);
    dfs(node->left,path,resCodeMap);
    path.pop_back();
    path.append(1);
    dfs(node->right,path,resCodeMap);
    path.pop_back();
}

void HuffmanCompression::generateCodeMap(unordered_map<unsigned char, hfmCodeBitSet> &resCodeMap)
{
    if(treeRoot == nullptr)
        return;
    hfmCodeBitSet path;
    dfs(treeRoot,path,resCodeMap);
}

void HuffmanCompression::generateEncode(const unsigned char *rawDataPtr, unsigned int rawDataSize,
                                        const unordered_map<unsigned char, hfmCodeBitSet> &resCodeMap,
                                        unsigned char* &outputDataPtr, unsigned int outputDataBitSize)
{
    outputDataPtr = new unsigned char[outputDataBitSize / 8 + 1];
    memset(outputDataPtr, 0, outputDataBitSize / 8 + 1);
    unsigned int outputBitOffset = 0;
    unsigned int codeBitLength = 0;
    unsigned int byteOff, bitOffInByte;

    // generate encode output data by looking up the weight map
    for ( unsigned int i = 0; i < rawDataSize; i++)
    {
        const hfmCodeBitSet& code = resCodeMap.at(rawDataPtr[i]);
        codeBitLength = code.getLength();
        for (unsigned int j = 0; j < codeBitLength; j++)
        {
            byteOff = (j + outputBitOffset) / 8;
            bitOffInByte = (j + outputBitOffset) % 8;
            outputDataPtr[byteOff] |= (code[j] << bitOffInByte);
        }
        outputBitOffset += codeBitLength;
    }
    if (outputBitOffset != outputDataBitSize)
    {
        ;
    }
}

unsigned int HuffmanCompression::calcEncodedOutputSize(const unordered_map<unsigned char, uint32_t> &weightMap,
                                                       const unordered_map<unsigned char, hfmCodeBitSet> &resCodeMap)
{
    unsigned int bitsCount = 0;
    unsigned char val;
    unsigned weight, codeLength;
    for(const auto &ele : weightMap)
    {
        val = ele.first;
        weight = ele.second;
        codeLength = static_cast<unsigned int>(resCodeMap.at(val).getLength());
        bitsCount += weight * codeLength;
    }
    return bitsCount;
}

void HuffmanCompression::getEncodedData(const unsigned char *rawDataPtr, unsigned int rawDataSize,
                                        unordered_map<unsigned char, unsigned int> &dstWeightMap,
                                        unsigned char * &outputDataPtr, unsigned int &outputDataBitSize)
{
    calcWeight(rawDataPtr,rawDataSize,dstWeightMap);
    generateEncodedNodeQueue(dstWeightMap);
    generateTree();
    unordered_map<unsigned char, hfmCodeBitSet> resCodeMap;
    generateCodeMap(resCodeMap);
    outputDataBitSize = calcEncodedOutputSize(dstWeightMap,resCodeMap);
    generateEncode(rawDataPtr, rawDataSize,resCodeMap,outputDataPtr,outputDataBitSize);
}

void HuffmanCompression::getDecodeData(const unsigned char *rawDataPtr, unsigned int rawDataBitSize,
                                       const vector <pair<unsigned char, unsigned int>> &srcWeightMapArr,
                                       unsigned char *&outputDataPtr, unsigned int &outputDataSize)
{
    generateDecodeNodeQueue(srcWeightMapArr);  // 根据读取文件之后构建的权值映射表构建节点队列
    generateTree();  // 建立编码树
    outputDataSize = calcDecodeOutputSize(srcWeightMapArr);  // 计算输出数据大小
    generateDecodedOutput(rawDataPtr, rawDataBitSize, outputDataPtr, outputDataSize);  // 对编码数据进行解码
}
void HuffmanCompression::generateDecodeNodeQueue(const vector<pair<unsigned char, unsigned int>> &weightMapArr)
{
    for( const auto &ele : weightMapArr )
    {
        nodeQueue.push(new TreeNode(ele.first,ele.second));
    }
}

unsigned int HuffmanCompression::calcDecodeOutputSize(const vector<pair<unsigned char, unsigned int>> &weightMap)
{
    unsigned int res = 0;
    for(const auto & ele : weightMap)
    {
        res+=ele.second;
    }
    return res;
}
void HuffmanCompression::generateDecodedOutput(const unsigned char *rawDataPtr, unsigned int rawDataBitSize,
                                               unsigned char *&outputDataPtr, unsigned int outputDataSize)
{
    outputDataPtr = new unsigned char[outputDataSize + 1];
    memset(outputDataPtr, 0, outputDataSize + 1);
    uint32_t bitOffInByte = 0;
    uint32_t byteOff = 0;
    uint32_t rawDataBitOff;
    uint32_t outputDataOff = 0;
    TreeNode *curNode = treeRoot;

    // generate decode output data by DFS the Huffman Tree
    for (rawDataBitOff = 0; rawDataBitOff < rawDataBitSize; rawDataBitOff++) {
        if (curNode -> left == nullptr && curNode -> right == nullptr) {
            outputDataPtr[outputDataOff++] = curNode -> val;
            curNode = treeRoot;
        }
        byteOff = rawDataBitOff / 8;
        bitOffInByte = rawDataBitOff % 8;
        if ((rawDataPtr[byteOff] >> bitOffInByte) & 1) {
            if (curNode == nullptr || curNode -> right == nullptr) {
                ;
            } else {
                curNode = curNode -> right;
            }
        } else {
            if (curNode == nullptr || curNode -> left == nullptr) {
                ;
            } else {
                curNode = curNode -> left;
            }
        }
    }
    if (curNode -> left == nullptr && curNode -> right == nullptr) {
        outputDataPtr[outputDataOff++] = curNode -> val;
        curNode = treeRoot;
    }
}